#!/usr/bin/python

## Alvar marker extraction script
## Created Jan 28, 2011
## Author: Sean Bowman, University of Minnesota

import sys
import os

# suppress alvar's writing to stderr when imported... maybe this is overkill.
devnull = open(os.devnull, 'w')
save_stderr = os.dup(2)
os.dup2(devnull.fileno(), 2)
import alvar
os.dup2(save_stderr, 2)

# The base name of the input files
# Files that fit the format base_file_name-num.[extension] will be analyzed
# by the script, where num is any number. Example: with base_file_name = "im"
# and extension = jpg, some valid images are im-3.jpg, im-01392.jpg, etc.

def usage():
	print "Usage: %s [startid] [endid] <options>" % sys.argv[0]
	print
	print "Mandatory arguments:"
	print "\t startid - the first image file ID to analyze"
	print "\t endid - the last image file ID to analyze"
	print
	print "Optional argument options, default in [ ]:"
	print "\t --ext=<type> \t\t the input image extension [pgm]"
	print "\t --out=<name> \t\t the output file name [landmarks.dat]"
	print "\t --indir=<dir> \t\t directory containing the input image files [.]"
	print "\t --outdir=<dir> \t\t the directory to write the output files [.]"
	print "\t --filename=<name> \t the base file name [image]"
	print "\t --marker-limit=<num> \t the max number of detectable markers [50]"
	print 
	
if __name__ == '__main__':
	if len(sys.argv) >= 2 and sys.argv[1] == '--help':
		usage(); exit(0)
		
	if len(sys.argv) < 3:
		print "Error: invalid arguments."
		print "Use --help for more information"
		exit(0)
	
	# parse required input arguments
	try:
		startid = int(sys.argv[1])
		endid = int(sys.argv[2])
	except ValueError:
		print "Invalid arguments."
		print "Use --help for more info."
		exit(1)
	
	#set default values
	base_file_name = "image"
	max_markers = 50 
	inputdir = '.'
	outputdir = '.'
	extension = 'pgm'
	output_file_name = 'landmarks.dat'
	
	# parse any included optional arguments
	for arg in sys.argv[3:]:
		flag, eq, val = arg.partition('=')
		if flag == '--ext':
			extension = val
		elif flag == '--out':
			output_file_name = val
		elif flag == '--indir':
			# strip a trailing '/' character if it is present
			if val[-1] == '/':
				val = val[:-1]
				
			inputdir = val
		elif flag == '--outdir':
			# strip a trailing '/' character if it is present
			if val[-1] == '/':
				val = val[:-1]
				
			outputdir = val
		elif flag == '--filename':
			base_file_name = val
		elif flag == '--max_markers':
			max_markers = int(val)
		elif flag == '--help':
			usage()
			exit(0)
		else:
			print "Invalid argument: %s" % str(arg)
			print "Use flag --help for more info."
			exit(1)
		
	alvar.init(max_markers) # TODO fix the memory leaks involved with this!
	
	# iterate over all valid files in the input directory
	# the input files must be sequential - it is not acceptable for image-4
	# and image-6 to be present without image-5, for example.
	directory = sorted(os.listdir(inputdir))
	filename_template = base_file_name + "-%d." + extension
		
	landmark_file = open(outputdir + '/' + output_file_name, 'w')
	
	for file_id in range(startid, endid + 1):
		# make sure that the current file exists
		file = filename_template % file_id
		if file not in directory:
			print "Error: file %d missing. Exiting." % file_id
			exit(1)
		
		num_markers = alvar.get_markers(inputdir + "/" + file)
			
		## keep a set of detected marker IDs so we don't get the same one twice,
		## as well as a set of "good" markers (ones detected only once) so we 
		## know which to write
		detected_ids = set()
		good_marker_ids = set()
		
		## iterate over all detected markers. we need to do this twice; once
		## to check for which markers are not duplicated, and once to write
		## those markers to the output files.
		for i in range(num_markers):
			marker_id = alvar.get_id(i)
			if marker_id in detected_ids:
				# duplicate!
				print "Marker %d is duplicated in %s, ignoring it" \
						% (marker_id, file)
				if marker_id in good_marker_ids:
					good_marker_ids.remove(marker_id)
			else:
				detected_ids.add(marker_id)
				good_marker_ids.add(marker_id)
		
		# make sure this image has detectable markers -- we don't want to
		# write a blank file.
		if len(good_marker_ids) == 0:
			continue
		
		## iterate over all detected markers
		for j in range(num_markers):
			marker_id = alvar.get_id(j)
			
			if marker_id not in good_marker_ids:
				# this marker was duplicated, ignore it
				continue
				
			landmark_file.write(str(marker_id) + " ")
			
			## iterate over all corners of this marker, printing coords
			for k in range(4):
				image_data = str(alvar.get_corners(j, k).x) + " " + \
							 str(alvar.get_corners(j, k).y)
				landmark_file.write(image_data + " ")
					
		landmark_file.write('\n')
